
package com.ctrip.corp.obt.shard.sql.mybatis.proxy;

import com.ctrip.corp.obt.generic.core.context.RequestContext;
import com.ctrip.corp.obt.generic.utils.Assert;
import com.ctrip.corp.obt.shard.context.PageContext;
import com.ctrip.corp.obt.shard.holder.PageHolder;
import com.ctrip.corp.obt.shard.sql.mybatis.interceptor.pagination.IPage;
import org.apache.ibatis.binding.BindingException;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.cursor.Cursor;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.StatementType;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.RowBounds;
import org.apache.ibatis.session.SqlSession;

import java.lang.reflect.Array;
import java.lang.reflect.Method;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;

import static com.ctrip.corp.obt.shard.context.PageContext.PAGE;

/**
 * @author marco.zheng
 * @version 1.0
 * @date 2022/03/17
 * override {@link MapperMethod}
 * quote {@link org.apache.ibatis.binding.MapperMethod.ParamMap}
 * quote {@link org.apache.ibatis.binding.MapperMethod.SqlCommand}
 * quote {@link org.apache.ibatis.binding.MapperMethod.MethodSignature}
 * add {@link OverrideMapperMethod#returnsPage()}
 */
public class OverrideMapperMethod {
    private final MapperMethod.SqlCommand command;
    private final MapperMethod.MethodSignature method;

    public OverrideMapperMethod(Class<?> mapperInterface, Method method, Configuration config) {
        this.command = new MapperMethod.SqlCommand(config, mapperInterface, method);
        this.method = new MapperMethod.MethodSignature(config, mapperInterface, method);
    }

    public Object execute(SqlSession sqlSession, Object[] args) {
        Object result;
        switch (command.getType()) {
            case INSERT: {
                // 这里获取的参数
                Object param = method.convertArgsToSqlCommandParam(args);
                // sqlSession.insert() 调用的写入sql的触发点
                result = rowCountResult(sqlSession.insert(command.getName(), param));
                // 如果这里在写入一次，会怎么，是否会操作两次sql，大概率是可以的，但是是否存在风险
                break;
            }
            case UPDATE: {
                Object param = method.convertArgsToSqlCommandParam(args);
                result = rowCountResult(sqlSession.update(command.getName(), param));
                break;
            }
            case DELETE: {
                Object param = method.convertArgsToSqlCommandParam(args);
                result = rowCountResult(sqlSession.delete(command.getName(), param));
                break;
            }
            case SELECT:
                if (method.returnsVoid() && method.hasResultHandler()) {
                    executeWithResultHandler(sqlSession, args);
                    result = null;
                    // unset page context if not returnsMany or returnsPage #Marco.Zheng
                    PageContext.unset();
                } else if (method.returnsMany()) {
                    result = executeForMany(sqlSession, args);
                } else if (method.returnsMap()) {
                    result = executeForMap(sqlSession, args);
                    // unset page context if not returnsMany or returnsPage #Marco.Zheng
                    PageContext.unset();
                } else if (method.returnsCursor()) {
                    result = executeForCursor(sqlSession, args);
                    // unset page context if not returnsMany or returnsPage #Marco.Zheng
                    PageContext.unset();
                } else if (returnsPage()) {
                    result = executeForIPage(sqlSession, args);
                } else {
                    Object param = method.convertArgsToSqlCommandParam(args);
                    result = sqlSession.selectOne(command.getName(), param);
                    if (method.returnsOptional()
                            && (result == null || !method.getReturnType().equals(result.getClass()))) {
                        result = Optional.ofNullable(result);
                    }
                    // unset page context if not returnsMany or returnsPage #Marco.Zheng
                    PageContext.unset();
                }
                break;
            case FLUSH:
                result = sqlSession.flushStatements();
                break;
            default:
                throw new BindingException("Unknown execution method for: " + command.getName());
        }
        if (result == null && method.getReturnType().isPrimitive() && !method.returnsVoid()) {
            throw new BindingException("Mapper method '" + command.getName()
                    + " attempted to return null from a method with a primitive return type (" + method.getReturnType() + ").");
        }
        return result;
    }

    /**
     * If return type is assignable from IPage or implement type @Marco.Zheng
     *
     * @return {@link IPage}
     */
    private boolean returnsPage() {
        return IPage.class.isAssignableFrom(method.getReturnType());
    }

    @SuppressWarnings("all")
    private <E> Object executeForIPage(SqlSession sqlSession, Object[] args) {
        IPage<E> result = null;
        if (Objects.nonNull(args)) {
            for (Object arg : args) {
                if (arg instanceof IPage) {
                    result = (IPage<E>) arg;
                    break;
                }
            }
        }
        // priority use the page from args #Marco.Zheng
        if (Objects.isNull(result)) {
            if (PageHolder.needPage(command.getName())) {
                result = PageContext.continueBuildPageWithPageQueryCondition(command.getName());
                Optional.ofNullable(result).ifPresent(PageContext::setPage);
            } else {
                PageContext.unset();
            }
        }
        Assert.notNull(result, "Can't found IPage for args!");
        Object param = method.convertArgsToSqlCommandParam(args);
        try {
            List<E> list = sqlSession.selectList(command.getName(), param);
            result.setRecords(list);
        } finally {
            Optional.ofNullable(RequestContext.getCurrentContext().getRequest()).ifPresent(request -> request.removeAttribute(PAGE));
            RequestContext.getCurrentContext().remove(PAGE);
        }
        return result;
    }

    private Object rowCountResult(int rowCount) {
        final Object result;
        if (method.returnsVoid()) {
            result = null;
        } else if (Integer.class.equals(method.getReturnType()) || Integer.TYPE.equals(method.getReturnType())) {
            result = rowCount;
        } else if (Long.class.equals(method.getReturnType()) || Long.TYPE.equals(method.getReturnType())) {
            result = (long) rowCount;
        } else if (Boolean.class.equals(method.getReturnType()) || Boolean.TYPE.equals(method.getReturnType())) {
            result = rowCount > 0;
        } else {
            throw new BindingException("Mapper method '" + command.getName() + "' has an unsupported return type: " + method.getReturnType());
        }
        return result;
    }

    private void executeWithResultHandler(SqlSession sqlSession, Object[] args) {
        MappedStatement ms = sqlSession.getConfiguration().getMappedStatement(command.getName());
        if (!StatementType.CALLABLE.equals(ms.getStatementType())
                && void.class.equals(ms.getResultMaps().get(0).getType())) {
            throw new BindingException("method " + command.getName()
                    + " needs either a @ResultMap annotation, a @ResultType annotation,"
                    + " or a resultType attribute in XML so a ResultHandler can be used as a parameter.");
        }
        Object param = method.convertArgsToSqlCommandParam(args);
        if (method.hasRowBounds()) {
            RowBounds rowBounds = method.extractRowBounds(args);
            sqlSession.select(command.getName(), param, rowBounds, method.extractResultHandler(args));
        } else {
            sqlSession.select(command.getName(), param, method.extractResultHandler(args));
        }
    }

    @SuppressWarnings("unchecked")
    private <E> Object executeForMany(SqlSession sqlSession, Object[] args) {
        IPage<E> page = null;
        // if need page, build in context #Marco.Zheng
        if (PageHolder.needPage(command.getName())) {
            page = PageContext.continueBuildPageWithPageQueryCondition(command.getName());
            PageContext.setPage(page);
            RequestContext.getCurrentContext().getRequest().setAttribute(PageContext.PAGE, PageContext.getPage());
        } else {
            PageContext.unset();
        }
        List<E> result;
        Object param = method.convertArgsToSqlCommandParam(args);

        try {
            if (method.hasRowBounds()) {
                RowBounds rowBounds = method.extractRowBounds(args);
                result = sqlSession.selectList(command.getName(), param, rowBounds);
            } else {
                result = sqlSession.selectList(command.getName(), param);
            }
            Optional.ofNullable(page).ifPresent(p -> p.setRecords(result));
        } finally {
            Optional.ofNullable(RequestContext.getCurrentContext().getRequest()).ifPresent(request -> request.removeAttribute(PAGE));
            RequestContext.getCurrentContext().remove(PAGE);
        }
        // issue #510 Collections & arrays support
        if (!method.getReturnType().isAssignableFrom(result.getClass())) {
            if (method.getReturnType().isArray()) {
                return convertToArray(result);
            } else {
                return convertToDeclaredCollection(sqlSession.getConfiguration(), result);
            }
        }
        return result;
    }

    private <T> Cursor<T> executeForCursor(SqlSession sqlSession, Object[] args) {
        Cursor<T> result;
        Object param = method.convertArgsToSqlCommandParam(args);
        if (method.hasRowBounds()) {
            RowBounds rowBounds = method.extractRowBounds(args);
            result = sqlSession.selectCursor(command.getName(), param, rowBounds);
        } else {
            result = sqlSession.selectCursor(command.getName(), param);
        }
        return result;
    }

    private <E> Object convertToDeclaredCollection(Configuration config, List<E> list) {
        Object collection = config.getObjectFactory().create(method.getReturnType());
        MetaObject metaObject = config.newMetaObject(collection);
        metaObject.addAll(list);
        return collection;
    }

    @SuppressWarnings("unchecked")
    private <E> Object convertToArray(List<E> list) {
        Class<?> arrayComponentType = method.getReturnType().getComponentType();
        Object array = Array.newInstance(arrayComponentType, list.size());
        if (arrayComponentType.isPrimitive()) {
            for (int i = 0; i < list.size(); i++) {
                Array.set(array, i, list.get(i));
            }
            return array;
        } else {
            return list.toArray((E[]) array);
        }
    }

    private <K, V> Map<K, V> executeForMap(SqlSession sqlSession, Object[] args) {
        Map<K, V> result;
        Object param = method.convertArgsToSqlCommandParam(args);
        if (method.hasRowBounds()) {
            RowBounds rowBounds = method.extractRowBounds(args);
            result = sqlSession.selectMap(command.getName(), param, method.getMapKey(), rowBounds);
        } else {
            result = sqlSession.selectMap(command.getName(), param, method.getMapKey());
        }
        return result;
    }
}
